import sys
sys.path.append('C:/Users/Hu/Dropbox/Research/PythonWork/Cancer/src/STAT/')

import numpy as np
import csv
from datetime import datetime
from time import strftime
import matplotlib.pyplot as plt
import math
import random
from scipy import stats
import ols

'''
Global and Local Empirical Bayes Smoothers with Gamma Model
'''

def getCurTime():
    """
    get current time
    Return value of the date string format(%Y-%m-%d %H:%M:%S)
    """
    format='%Y-%m-%d %H:%M:%S'
    sdate = None
    cdate = datetime.now()
    try:
        sdate = cdate.strftime(format)
    except:
        raise ValueError
    return sdate

def build_data_list(inputCSV):
    sKey = []
    fn = inputCSV
    f = open(inputCSV)
    #ra = csv.DictReader(file(fn), dialect="excel")
    ra = csv.DictReader(f, dialect="excel")
    
    for record in ra:
        #print record[ra.fieldnames[0]], type(record[ra.fieldnames[-1]])
        for item in ra.fieldnames:
            temp = float(record[item])
            sKey.append(temp)
    sKey = np.array(sKey)
    sKey.shape=(-1,len(ra.fieldnames))
    return sKey

def log(data):
    logdata = []
    for item in data:
        logdata.append(math.log(item))
    return logdata

def pickBreakpoint(response, predictor):
    bpChoices = np.sort(predictor)
    results = np.zeros((len(predictor)-1, 2))
    
    for i in range(len(predictor)-1):
        x2star = (predictor - bpChoices[i]) * np.greater(predictor, bpChoices[i])   
        tempPredictor = np.array(zip(predictor, x2star))
        tempmodel = ols.ols(response, tempPredictor,'y',['x1', 'x2'])
        results[i,0] = i
        results[i,1] = tempmodel.sse

    optBP = int(results[np.argmin(results, axis = 0)[1],0])
    print optBP, 'Optimal changepoint: ', bpChoices[optBP], ' with SSE = ', results[optBP, 1]

    x2star = (predictor - bpChoices[optBP]) * np.greater(predictor, bpChoices[optBP])
    optPredictor = np.array(zip(predictor, x2star))
    optmodel = ols.ols(response, optPredictor,'y',['x1', 'x2'])
    
    return bpChoices[optBP], results, optmodel, optmodel.b[0]+optmodel.b[1]*predictor+optmodel.b[2]*x2star

#--------------------------------------------------------------------------
#MAIN

if __name__ == "__main__":
    print '===================================================='
    print "begin at " + getCurTime()
    filePath = 'C:/_DATA/migration_census_2000/Aggregation/'
    file = filePath + 'census_county_migration_aggregation_dist_50k_bin.csv'
    data = build_data_list(file)

    file = filePath + 'census_county_migration_aggregation_expFlow_dist_50k_bin.csv'
    expdata50 = build_data_list(file)
  
    avgExp50 = np.average(expdata50[:,1])
    avgOb = np.average(data[:,1])

    ajustOB = []
    for ob, exp in zip(data[:,1], expdata50[:,1]):
        ajustOB.append(ob/(exp/avgExp50))

    print sum(ajustOB), sum(data[:,1])
        
    expmodel50 = ols.ols(expdata50[:,1],expdata50[:,0],'y',['x'])
    #print mymodel.summary()

    #y = np.log(ajustOB)
    #x = np.log(data[:,0])
    
    #slope, intercept, r_value, p_value, std_err = stats.linregress(x,y)
    #print('Linear regression using stats.linregress')
    #print('regression: b0 = %.4f, b1 = %.4f, r_square = %.4f, p_value = %.4f, std error = %.4f' % (intercept,slope,r_value**2, p_value,stderr))
    '''
    plt.subplot(211)
    plt.scatter(expdata20[:,0]/1000, expdata20[:,1], label = 'Original data', color ='b')
    #plt.plot(expdata20[:,0], expmodel20.b[1]*expdata20[:,0] + expmodel20.b[0], 'b', label='Fitted line', linewidth=1.5)
    #plt.plot(x, np.average(y), 'g', label='Average line')
    plt.axhline(y=avgExp20, color ='b', label='Average line', linewidth=1.5)
    #plt.plot(data[:,0], data[:,1], 'k--')
    plt.title('Census Data (bin = 20 km)')
    plt.xlabel('distance (km)')
    plt.ylabel('exp flow volumn')
    plt.axis([min(expdata20[:,0]/1000) * 0.9, max(expdata20[:,0]/1000) *1.1, min(expdata20[:,1]) * 0.9, max(expdata20[:,1]) * 1.1])
    plt.legend()

    plt.subplot(212)
    plt.scatter(expdata50[:,0]/1000, expdata50[:,1], label = 'Original data', color ='g')
    #plt.plot(expdata50[:,0], expmodel50.b[1]*expdata50[:,0] + expmodel50.b[0], 'g', label='Fitted line', linewidth=1.5)
    #plt.plot(x, np.average(y), 'g', label='Average line')
    plt.axhline(y=avgExp50, color ='g', label='Average line', linewidth=1.5)
    #plt.plot(data[:,0], data[:,1], 'k--')
    plt.title('Census Data (bin = 50 km)')
    plt.xlabel('distance (km)')
    plt.ylabel('exp flow volumn')
    plt.axis([min(expdata50[:,0]/1000) * 0.9, max(expdata50[:,0]/1000) *1.1, min(expdata50[:,1]) * 0.9, max(expdata50[:,1]) * 1.1])
    plt.legend()
    
    plt.show()
    '''

    avgExp50 = np.average(expdata50[:,1])


    ajustOB50 = []
    for ob, exp in zip(data[:,1], expdata50[:,1]):
        ajustOB50.append(ob/(exp/avgExp50))

    print sum(ajustOB50)

    adjmodel50 = ols.ols(np.log(ajustOB50), np.log(data[:,0]),'y',['x'])
    obmodel = ols.ols(np.log(data[:,1]), np.log(data[:,0]),'y',['x'])

    optBP_adj50, results_adj50, optmodel_adj50, y_hat_adj50 = pickBreakpoint(np.log(ajustOB50), np.log(data[:,0]))
    optBP_ob, results_ob, optmodel_ob, y_hat_ob = pickBreakpoint(np.log(data[:,1]), np.log(data[:,0]))
    
    print 'Simple Linear Regression for Original Data.............'
    print obmodel.summary()
    print 'Piecewise Regression for Original Data.............'
    print optmodel_ob.summary()
    print 'Simple Linear Regression for Adjusted Data.............'
    print adjmodel50.summary()
    print 'Piecewise Regression for Adjusted Data.............'
    print optmodel_adj50.summary()

    f = plt.figure()
    f.text(0.5,0.975,'Census Data (bin = 50 km)', horizontalalignment='center',verticalalignment='top', fontsize=20)

    
    '''
    plt.subplot(211)
    plt.scatter(data[:,0], data[:,1], label = 'Original data', color ='b')
    plt.scatter(data[:,0], ajustOB50, label = 'Adjusted data', color ='g')
    plt.axis([min(data[:,0]) * 0.9, max(data[:,0]) *1.1, min(data[:,1]) * 0.9, max(data[:,1]) * 1.1])
    plt.xlabel('distance (meter)')
    plt.ylabel('flow volumn')
    plt.legend()
    '''
    

    #plt.subplot(212)
    plt.loglog(data[:,0], data[:,1], marker = 'o', linestyle = '', label = 'Original data', ms=4, color ='b')
    plt.loglog(data[:,0], np.exp(obmodel.b[1]*np.log(data[:,0]) + obmodel.b[0]), 'b', linestyle ='--', label='Simple Linear Regression for original data', linewidth=2)
    plt.loglog(data[:,0], np.exp(y_hat_ob), 'b', label='Piecewise Regression for original data', linewidth=2)
    plt.loglog(data[:,0], ajustOB50, marker = '^', linestyle = '', label = 'Adjusted data', ms=6, color ='g')
    plt.loglog(data[:,0], np.exp(adjmodel50.b[1]*np.log(data[:,0]) + adjmodel50.b[0]), linestyle ='--', color ='g', label='Simple Linear Regression for adjusted data', linewidth=2)
    plt.loglog(data[:,0], np.exp(y_hat_adj50), 'g', label='Piecewise Regression for adjusted data', linewidth=2)
    
    #plt.semilogx(np.log(data[:83,0]), obmodel.b[1]*np.log(data[:83,0]) + obmodel.b[0], 'r')
    plt.axis([min(data[:,0]) * 0.8, max(data[:,0]) *1.5, min(data[:83,1]) * 0.8, max(data[:,1]) * 3.5])
    
    plt.tick_params(which = 'major', size = 14, width = 2)
    plt.tick_params(which = 'minor', size = 8, width = 1)
    
    plt.xlabel('distance (km)')
    plt.ylabel('flow volumn')
    plt.legend(loc=3)
    plt.show()

    #fileLoc = filePath + 'segmented_data.csv'
    #np.savetxt(fileLoc, zip(np.log(ajustOB50), np.log(data[:,0])), delimiter=',', fmt = '%s')
    
